package com.zws.cucumber.util.cass;

import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.cql.*;
import com.datastax.oss.driver.api.core.metadata.schema.RelationMetadata;
import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
import com.datastax.oss.driver.api.core.session.Session;
import com.datastax.oss.driver.api.core.type.*;
import com.zws.cucumber.adaptation.cass.ZwsCqlSessionProxy;
import com.zws.cucumber.bootstrap.ZwsCassBootstrap;
import com.zws.cucumber.util.colorfullog.ZwsPrintTableHelper;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static com.datastax.oss.driver.api.core.type.DataTypes.*;
import static com.zws.cucumber.adaptation.cass.ZwsCqlSessionProxy.markModifiedTables;

public class ZwsCassTemplateImpl implements ZwsCassTemplate {

    private static final Logger logger = LoggerFactory.getLogger(ZwsCassTemplateImpl.class);
    private static final ConcurrentMap<String, String> modifiedTables = new ConcurrentHashMap<>();
    private static volatile Collection<String> cachedTables;

    private static volatile Set<String> views;
    private static volatile Map<String, String> view2BaseTableMapping;
    protected CqlSession session;

    public ZwsCassTemplateImpl(CqlSession session) {
        this.session = session;
    }

    public static void addModifiedTable(String table) {
        if (StringUtils.isNotBlank(table)) {
            int idx = table.indexOf('.');
            if (idx >= 0) {
                table = table.substring(idx + 1).toLowerCase();
            }
            modifiedTables.put(table, table);
        }
    }

    public static void addModifiedTables(Iterable<String> tables) {
        for (String table : tables) {
            if (StringUtils.isNotBlank(table)) {
                int idx = table.indexOf('.');
                if (idx >= 0) {
                    table = table.substring(idx + 1).toLowerCase();
                }
                modifiedTables.put(table, table);
            }
        }
    }

    public static boolean hasModifiedTable(Session session, String table) {
        if (ZwsCassTemplateImpl.views == null || ZwsCassTemplateImpl.view2BaseTableMapping == null) {
            // init
            Set<String> views = new HashSet<>(256);
            Map<String, String> view2BaseTableMapping = new HashMap<>(256);
            session.getMetadata().getKeyspace(session.getKeyspace().get()).get().getViews().values()
                    .stream()
                    .forEach(v -> {
                        views.add(v.getName().asInternal());
                        view2BaseTableMapping.put(v.getName().asInternal(), v.getBaseTable().asInternal());
                    });
            ZwsCassTemplateImpl.views = views;
            ZwsCassTemplateImpl.view2BaseTableMapping = view2BaseTableMapping;
        }
        return modifiedTables.containsKey(table)
                || (views.contains(table) && modifiedTables.containsKey(view2BaseTableMapping.get(table)));
    }

    public static void setCachedTables(List<String> cachedTables) {
        ZwsCassTemplateImpl.cachedTables = cachedTables;
    }

    public static Map<String, String> getModifiedTables() {
        return new HashMap<>(modifiedTables);
    }

    public static void putAllModifiedTables(Map<String, String> modifiedTables) {
        ZwsCassTemplateImpl.modifiedTables.putAll(modifiedTables);
    }

    @Override
    public List<Map<String, Object>> queryList(String cql) {
        logger.debug("Executing SQL statement [" + cql + "]");
        List<Map<String, Object>> results = new ArrayList<Map<String, Object>>();
        ResultSet resultSet = execute0(cql);
        if (resultSet.getAvailableWithoutFetching() == 0) {
            logger.debug("result empty!");
            return results;
        }
        ColumnDefinitions definitions = resultSet.getColumnDefinitions();

        List<Row> rows = resultSet.all();
        for (Row row : rows) {
            Map<String, Object> result = convertToMap(definitions, row);
            results.add(result);
        }

        return results;
    }

    @Override
    public Pair<Map<String, ColumnMetadata>, List<Map<String, Object>>> queryListWithMetadata(String cql) {
        logger.debug("Executing SQL statement [" + cql + "]");
        List<Map<String, Object>> results = new ArrayList<Map<String, Object>>();
        ResultSet resultSet = execute0(cql);

        ColumnDefinitions definitions = resultSet.getColumnDefinitions();
        List<Row> rows = resultSet.all();
//        if (rows.size() == 0) {
//            logger.debug("result empty!");
//            return Pair.of(new LinkedHashMap<>(), new ArrayList<>());
//        }
        for (Row row : rows) {
            Map<String, Object> result = convertToMap(definitions, row);
            results.add(result);
        }

        Map<String, ColumnMetadata> metadataMap = new LinkedHashMap<>();
        for (ColumnDefinition def : definitions) {
            String name = def.getName().asInternal().toLowerCase();
            metadataMap.put(name, new ColumnMetadata(name, convertType(def.getType())));
        }

        return Pair.of(metadataMap, results);
    }

    private Map<String, Object> convertToMap(ColumnDefinitions definitions, Row row) {
        Map<String, Object> result = new HashMap<>();
        definitions.forEach(x -> {
            String name = x.getName().toString();
            String key = name.toUpperCase();
            DataType type = x.getType();

            if (type.equals(TEXT)) {
                result.put(key, row.getString(name));
            } else if (type.equals(DataTypes.TIMESTAMP)) {
                result.put(key, row.getLocalTime(name));
            } else if (type.equals(DataTypes.VARINT)) {
                result.put(key, row.getLong(name));
            } else if (type.equals(DataTypes.DOUBLE)) {
                result.put(key, row.isNull(name) ? null : row.getDouble(name));
            } else if (type.equals(DataTypes.FLOAT)) {
                result.put(key, row.isNull(name) ? null : row.getFloat(name));
            } else if (type.equals(DataTypes.BIGINT)) {
                result.put(key, row.isNull(name) ? null : row.getLong(name));
            } else if (type.equals(DataTypes.INT)) {
                result.put(key, row.isNull(name) ? null : row.getInt(name));
            } else if (type.equals(DataTypes.BOOLEAN)) {
                result.put(key, row.isNull(name) ? null : row.getBoolean(name));
            } else if (type.equals(DataTypes.SMALLINT)) {
                result.put(key, row.isNull(name) ? null : row.getShort(name));
            } else {
                result.put(key, row.getObject(name));
            }
        });
        return result;
    }

    @Override
    public void execute(String cql) {
        logger.debug("Executing SQL statement [" + cql + "]");
        ResultSet resultSet = execute0(cql);
        logger.debug("execute end: " + (resultSet.one() == null ?
                "null" : resultSet.one().getFormattedContents()));
    }

    @Override
    public void executeBatch(List<String> cqlList) {
        logger.debug("Execute Batch Statement [{}]", cqlList
                .stream()
                .collect(Collectors.joining()));

        long s = System.currentTimeMillis();
        BatchStatement batchStatement = BatchStatement.newInstance(BatchType.LOGGED);
        for (String cql : cqlList) {
            SimpleStatement simpleStatement = SimpleStatement.newInstance(cql);
            if (batchStatement.computeSizeInBytes(session.getContext()) + simpleStatement.computeSizeInBytes(session.getContext()) + 512 > 51200) {
                execute0(batchStatement);
                batchStatement = BatchStatement.newInstance(BatchType.LOGGED);
            }
            batchStatement = batchStatement.add(simpleStatement);
        }
        if (batchStatement.size() > 0) {
            execute0(batchStatement);
        }
        long e = System.currentTimeMillis();
        logger.debug("execute done, elapsed: {}ms", e - s);
    }

    @Override
    public Map<String, ColumnMetadata> queryMetadata(String tableName) {
        String keyspace = tableName.contains(".") ? tableName.substring(0, tableName.indexOf('.')) : "com/zws";
        tableName = tableName.substring(tableName.indexOf('.') + 1);
        TableMetadata tableMetadata = session
                .getMetadata()
                .getKeyspace(session.getKeyspace().get()).get()
                .getTable(tableName).get();

        Map<String, ColumnMetadata> metadataMap = new LinkedHashMap<>();
        tableMetadata.getColumns().forEach((k, v) -> {
            metadataMap.put(k.asInternal().toLowerCase(),
                    new ColumnMetadata(k.asInternal().toLowerCase(), convertType(v.getType())));
        });
        return metadataMap;
    }

    private ColumnType convertType(DataType type) {
        if (type == TINYINT
                || type == SMALLINT
                || type == INT
                || type == BIGINT
                || type == DOUBLE
                || type == FLOAT
                || type == DECIMAL) {
            return ColumnType.INT;
        }
        if (type == TEXT) {
            return ColumnType.STR;
        }
        if (type instanceof SetType) {
            return ColumnType.SET;
        }
        if (type instanceof ListType) {
            return ColumnType.LIST;
        }
        if (type instanceof MapType) {
            return ColumnType.MAP;
        }
        if (type instanceof UserDefinedType) {
            return ColumnType.UDT;
        }
        return ColumnType.NIL;
    }

    public List<String> getTables(String keyspace) {
        return session.getMetadata().getKeyspace(keyspace).get()
                .getTables()
                .values()
                .stream()
                .map(RelationMetadata::getName)
                .map(Object::toString)
                .collect(Collectors.toList());
    }

    @Override
    public void cleanData(boolean onlyModifiedTables, String keyspace) {
        Collection<String> tables;
        if (onlyModifiedTables) {
            tables = modifiedTables.keySet();
        } else {
            tables = cachedTables;
            if (tables == null) {
                tables = session.getMetadata().getKeyspace(keyspace).get().getTables().values()
                        .stream()
                        .map(RelationMetadata::getName)
                        .map(Object::toString)
                        .map(s -> keyspace + "." + s)
                        .collect(Collectors.toList());
                cachedTables = tables;
            }
        }
        try {
            if (!tables.isEmpty()) {
                if (ZwsCassBootstrap.cassBackdoorClient != null) {
                    tables = tables
                            .stream()
                            .map(s -> s.substring(s.indexOf('.') + 1))
                            .collect(Collectors.toList());
                    ZwsCassBootstrap.cassBackdoorClient.truncate(keyspace, tables.toArray(new String[0]));
                } else {
                    truncate(tables);
                }
            } else {
                logger.debug("modified tables is empty, no need to truncate");
            }
        } finally {
            if (onlyModifiedTables) {
                modifiedTables.clear();
            }
        }
    }

    @Override
    public void cleanData(String keyspace) {
        cleanData(true, keyspace);
    }

    @Override
    public void cleanData(String keyspace, String... tables) {
        if (tables.length > 0) {
            if (ZwsCassBootstrap.cassBackdoorClient != null) {
                tables = Stream.of(tables)
                        .map(s -> s.substring(s.indexOf('.') + 1))
                        .toArray(String[]::new);
                ZwsCassBootstrap.cassBackdoorClient.truncate(keyspace, tables);
            } else {
                truncate(tables);
            }
        } else {
            logger.debug("tables is empty, no need to truncate");
        }
    }

    @Override
    public void truncate(String... tables) {
        truncate(Arrays.asList(tables));
    }

    /*
     * Removes all data from the specified table immediately and irreversibly, and removes all data from any materialized views derived from that table.
     * https://docs.datastax.com/en/cql-oss/3.x/cql/cql_reference/cqlTruncate.html
     */
    private void truncate(Collection<String> tables) {
        logger.debug("truncated table: {}", tables);

        String cqlPattern = "truncate table %s";
        List<CompletionStage<AsyncResultSet>> stages = new ArrayList<>();
        tables.forEach(tableName -> {
            stages.add(session.executeAsync(String.format(cqlPattern, tableName)));
        });

        stages.forEach(stage -> {
            try {
                stage.toCompletableFuture().get();
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
    }

    @Override
    public void printTable(String tableName) {
        printTable(tableName, null);
    }

    @Override
    public void printTable(String tableName, String whereClause) {
        String cql;
        if (StringUtils.isNotBlank(whereClause)) {
            cql = String.join(" ", "select", "*", "from", tableName, "where", whereClause);
        } else {
            cql = String.join(" ", "select", "*", "from", tableName);
        }
        logger.debug("cql: {}", cql);

        Pair<Map<String, ColumnMetadata>, List<Map<String, Object>>> rawPair = queryListWithMetadata(cql);
        Pair<Map<String, ColumnMetadata>, List<Map<String, Object>>> pair = Pair.of(rawPair.getKey(), toLowerCase(rawPair.getValue()));
        logger.trace("query result: {}", pair);

        ZwsPrintTableHelper.printTable0(
                tableName, Pair.of(new ArrayList<>(pair.getLeft().keySet()), pair.getRight()));
    }

    private <V> List<Map<String, V>> toLowerCase(List<Map<String, V>> originList) {
        List<Map<String, V>> resultList = new ArrayList<>();
        for (Map<String, V> origin : originList) {
            Map<String, V> result = new LinkedHashMap<>();
            origin.forEach((k, v) -> result.put(k.toLowerCase(), v));
            resultList.add(result);
        }
        return resultList;
    }

    private <T> void preprocess(T stmt) {
        ZwsCqlSessionProxy.waitAsyncExecution();
//        ZwsCqlSessionProxyHelper.waitAsyncCleanData();
        markModifiedTables(stmt);
//        ZwsCqlSessionProxyHelper.intercept(stmt);
    }

    private ResultSet execute0(String cql) {
        preprocess(cql);
        return session.execute(cql);
    }

    private ResultSet execute0(Statement<?> stmt) {
        preprocess(stmt);
        return session.execute(stmt);
    }
}
